#!/usr/bin/env python3
"""
读取 ROS2 点云和 IMU 数据，显示：
1. IMU数据积分还原世界系轨迹（红色线）
2. 点云数据用对应时刻IMU姿态/位置还原世界系坐标，按intensity用冷色渐变显示
3. 标题显示总IMU轨迹数量及点云数据数量
Usage:
    python3 plot_imu_cloud_intensity.py [--bag xxx.db3]
"""
import argparse
import numpy as np
import matplotlib.pyplot as plt
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions, StorageFilter
from sensor_msgs.msg import PointCloud2, Imu
from rclpy.serialization import deserialize_message
G = 9.80665  # gravity magnitude used for checks (m/s^2)
# 重力加速度
GRAVITY = np.array([0.0, 0.0, 9.66723215623791], dtype=np.float32)
# GRAVITY = np.array([0.0, 0.0, G], dtype=np.float32)

# 四元数转旋转矩阵
def quat2R(qw, qx, qy, qz):
    R = np.array([
        [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
        [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)],
        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]
    ], dtype=np.float32)
    return R

# PointCloud2 转 xyz/intensity

def pc2_to_xyz_intensity(pc: PointCloud2):
    arr = np.frombuffer(pc.data, np.uint8).reshape(-1, pc.point_step)
    xyz = arr[:, :12].copy().view(np.float32).reshape(-1, 3)
    # intensity 通常在第16字节（float32），具体视fields而定
    intensity = None
    for f in pc.fields:
        if f.name == 'intensity':
            offset = f.offset
            intensity = arr[:, offset:offset+4].copy().view(np.float32).reshape(-1)
            break
    if intensity is None:
        intensity = np.zeros(xyz.shape[0], dtype=np.float32)
    mask = np.isfinite(xyz).all(axis=1)
    return xyz[mask], intensity[mask]

# 加载全部数据
def load_data(bag_path, max_seconds=None):
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'), ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=['/unilidar/imu', '/unilidar/cloud']))
    latest_R = np.eye(3, dtype=np.float32)
    latest_v = np.zeros(3, dtype=np.float32)
    latest_t = np.zeros(3, dtype=np.float32)
    last_time = None
    imu_traj = []
    cloud_pts = []
    cloud_intensity = []
    imu_times = []
    start_time = None
    while reader.has_next():
        topic, data, t = reader.read_next()
        if topic == '/unilidar/imu':
            imu = deserialize_message(data, Imu)
            curr_time = imu.header.stamp.sec + imu.header.stamp.nanosec * 1e-9 if hasattr(imu, 'header') else t * 1e-9
            if start_time is None:
                start_time = curr_time
            if max_seconds is not None and curr_time - start_time > max_seconds:
                break
            if last_time is None:
                dt = 0.01
            else:
                dt = curr_time - last_time
            last_time = curr_time
            q = imu.orientation
            latest_R = quat2R(q.w, q.x, q.y, q.z)
            acc_body = np.array([imu.linear_acceleration.x,
                                 imu.linear_acceleration.y,
                                 imu.linear_acceleration.z], dtype=np.float32)
            acc_world = latest_R @ acc_body - GRAVITY
            # 将Z轴加速度中与重力加速度接近的部分置零，减少漂移
            # if np.abs(acc_world[2] + G) < 1.0:
            # acc_world[2] = 0.0
            latest_v += acc_world * dt
            latest_t += latest_v * dt
            imu_traj.append(latest_t.copy())
            imu_times.append(curr_time - start_time)
        elif topic == '/unilidar/cloud':
            if max_seconds is not None and last_time is not None and last_time - start_time > max_seconds:
                break
            pc = deserialize_message(data, PointCloud2)
            pts_imu, intensity = pc2_to_xyz_intensity(pc)
            pts_world = (latest_R @ pts_imu.T).T + latest_t
            cloud_pts.append(pts_world)
            cloud_intensity.append(intensity)
    imu_traj = np.vstack(imu_traj) if imu_traj else np.empty((0, 3))
    cloud_pts = np.vstack(cloud_pts) if cloud_pts else np.empty((0, 3))
    cloud_intensity = np.concatenate(cloud_intensity) if cloud_intensity else np.empty((0,))
    imu_times = np.array(imu_times)
    return imu_traj, cloud_pts, cloud_intensity, imu_times

# 可视化
def plot_all(imu_traj, cloud_pts, cloud_intensity, imu_times, show_pcd_rate, show_traj_only):
    fig = plt.figure(figsize=(14, 6))
    ax1 = fig.add_subplot(111, projection='3d')
    ax1.set_xlabel('X world')
    ax1.set_ylabel('Y world')
    ax1.set_zlabel('Z world')
    # IMU轨迹
    if imu_traj.shape[0] > 1:
        ax1.plot(imu_traj[:, 0], imu_traj[:, 1], imu_traj[:, 2], c='red', label='IMU traj')
    if not show_traj_only:
        if show_pcd_rate is not None and 0 < show_pcd_rate < 1.0:
            sample_size = int(cloud_pts.shape[0] * show_pcd_rate)
            if sample_size > 0:
                indices = np.random.choice(cloud_pts.shape[0], size=sample_size, replace=False)
                cloud_pts = cloud_pts[indices]
                cloud_intensity = cloud_intensity[indices]
        # 点云，按intensity用冷色渐变
        if cloud_pts.shape[0] > 0:
            norm_intensity = (cloud_intensity - cloud_intensity.min()) / (np.ptp(cloud_intensity) + 1e-6)
            colors = plt.cm.cool(norm_intensity)
            ax1.scatter(cloud_pts[:, 0], cloud_pts[:, 1], cloud_pts[:, 2], s=1, c=colors, label='Cloud (intensity)')
        ax1.legend(loc='upper right')
        ax1.set_title(f'IMU轨迹数:{imu_traj.shape[0]} 点云数:{cloud_pts.shape[0]}')

    # plt.tight_layout()
    plt.show(block=True)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', default='/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3', help='ros2 bag db3 path')
    parser.add_argument('--seconds', type=float, default=0.0, help='只读取前N秒数据，默认60s，空则全部')
    parser.add_argument('--show_pcd_rate', type=float, default=0.001, help='点云显示采样率，默认0.01即1%')
    parser.add_argument('--show_traj_only', type=bool, default=True, help='仅显示轨迹')
    args = parser.parse_args()
    print('读取 bag 并加载 IMU/点云 世界系轨迹 …')
    max_seconds = args.seconds if args.seconds > 0 else None
    imu_traj, cloud_pts, cloud_intensity, imu_times = load_data(args.bag, max_seconds)
    print(f'IMU轨迹数:{imu_traj.shape[0]} 点云数:{cloud_pts.shape[0]}')
    plot_all(imu_traj, cloud_pts, cloud_intensity, imu_times, args.show_pcd_rate, args.show_traj_only)
